//
//  XWCustomURLProtocol.m
//  DNS拦截
//
//  Created by Xie Wei on 2019/7/29.
//  Copyright © 2019 Xie Wei. All rights reserved.
//

#import "XWCustomURLProtocol.h"
#import "XWSessionConfiguration.h"
#import <HappyDNS/HappyDNS.h>
#import "XWDNSModel.h"

static NSString * const URLProtocolHandledKey = @"URLProtocolHandledKey";

@interface XWCustomURLProtocol ()<NSURLSessionDelegate>

@property (nonatomic, strong) NSURLSession *session;

@end

@implementation XWCustomURLProtocol

// 开始监听
+ (void)startMonitor {
    XWSessionConfiguration *sessionConfiguration = [XWSessionConfiguration defaultConfiguration];
    [NSURLProtocol registerClass:[XWCustomURLProtocol class]];
    if (![sessionConfiguration isExchanged]) {
        [sessionConfiguration load];
    }
}

// 停止监听
+ (void)stopMonitor {
    XWSessionConfiguration *sessionConfiguration = [XWSessionConfiguration defaultConfiguration];
    [NSURLProtocol unregisterClass:[XWCustomURLProtocol class]];
    if ([sessionConfiguration isExchanged]) {
        [sessionConfiguration unload];
    }
}


+ (BOOL)canInitWithRequest:(NSURLRequest *)request {
    NSString * scheme = [[request.URL scheme] lowercaseString];
    
    // 看看是否已经处理过了，防止无限循环 根据业务来截取
    if ([NSURLProtocol propertyForKey: URLProtocolHandledKey inRequest:request]) {
        return NO;
    }
    
    // TODO - 这里是自己需要处理的逻辑
    if ([scheme isEqual:@"http"] || [scheme isEqual:@"https"]) {
        return YES;
    }
    
    return NO;
}

+ (NSURLRequest *)canonicalRequestForRequest:(NSURLRequest *)request {
    return request;
}

+ (BOOL)requestIsCacheEquivalent:(NSURLRequest *)a toRequest:(NSURLRequest *)b {
    return [super requestIsCacheEquivalent:a toRequest:b];
}

- (id)initWithRequest:(NSURLRequest *)request cachedResponse:(NSCachedURLResponse *)cachedResponse client:(id <NSURLProtocolClient>)client {
    return [super initWithRequest:request cachedResponse:cachedResponse client:client];
}

//开始请求
- (void)startLoading
{
    
    NSLog(@"***监听接口：%@", self.request.URL.absoluteString);
    
    NSMutableURLRequest *mutableReqeust = [[self request] mutableCopy];
    //标示该request已经处理过了，防止无限循环
    [NSURLProtocol setProperty:@(YES) forKey:URLProtocolHandledKey inRequest:mutableReqeust];
    
    // dns解析
    NSMutableURLRequest *request = [self.class replaceHostInRequset:mutableReqeust];
    
    // 使用NSURLSession继续把request发送出去
    NSURLSessionConfiguration *config = [NSURLSessionConfiguration defaultSessionConfiguration];
    NSOperationQueue *mainQueue = [NSOperationQueue mainQueue];
    self.session = [NSURLSession sessionWithConfiguration:config delegate:self delegateQueue:mainQueue];
    NSURLSessionDataTask *task = [self.session dataTaskWithRequest:request];
    [task resume];
}

+ (NSMutableURLRequest *)replaceHostInRequset:(NSMutableURLRequest *)request {
    NSLog(@"%s", __func__);
    if ([request.URL host].length == 0) {
        return request;
    }
    
    NSString *originUrlString = [request.URL absoluteString];
    NSString *originHostString = [request.URL host];
    NSRange hostRange = [originUrlString rangeOfString:originHostString];
    if (hostRange.location == NSNotFound) {
        return request;
    }
    
    // 用HappyDNS 替换host
    NSMutableArray *array = [NSMutableArray array];
    // 第一dns解析为114，第二解析才是系统dns
    NSString *dns = [XWDNSModel shared].dns;
    [array addObject:[[QNResolver alloc] initWithAddress:dns]];
    [array addObject:[QNResolver systemResolver]];
    QNDnsManager *dnsManager = [[QNDnsManager alloc] init:array networkInfo:[QNNetworkInfo normal]];
    NSArray *queryArray = [dnsManager query:originHostString];
    if (queryArray && queryArray.count > 0) {
        NSString *ip = queryArray[0];
        if (ip && ip.length) {
            // 替换host
            NSString *urlString = [originUrlString stringByReplacingCharactersInRange:hostRange withString:ip];
            NSURL *url = [NSURL URLWithString:urlString];
            request.URL = url;
            
            [request setValue:originHostString forHTTPHeaderField:@"Host"];
        }
    }
    
    return request;
}

//结束请求
- (void)stopLoading {
    [self.session invalidateAndCancel];
    self.session = nil;
}

#pragma mark - NSURLSessionDelegate
- (void)URLSession:(NSURLSession *)session dataTask:(NSURLSessionDataTask *)dataTask didReceiveResponse:(NSURLResponse *)response completionHandler:(void (^)(NSURLSessionResponseDisposition))completionHandler
{
    [self.client URLProtocol:self didReceiveResponse:response cacheStoragePolicy:NSURLCacheStorageNotAllowed];
    completionHandler(NSURLSessionResponseAllow);
}

- (void)URLSession:(NSURLSession *)session dataTask:(NSURLSessionDataTask *)dataTask didReceiveData:(NSData *)data
{
    // 打印返回数据
    NSString *dataStr = [[NSString alloc] initWithData:data encoding:NSUTF8StringEncoding];
    if (dataStr) {
        NSLog(@"***截取数据***: %@", dataStr);
    }
    [self.client URLProtocol:self didLoadData:data];
}

- (void)URLSession:(NSURLSession *)session task:(NSURLSessionTask *)task didCompleteWithError:(NSError *)error
{
    if (error) {
        [self.client URLProtocol:self didFailWithError:error];
    } else {
        [self.client URLProtocolDidFinishLoading:self];
    }
}


@end
